# simpleNetwork
# Copyright (c) 2025 charlieshu <https://gitee.com/charlieshu>
# Released under the MIT License.
# See LICENSE file for full text.

import socket
import threading
import time
import os
from datetime import datetime
from typing import List, Dict

# =================== 全局变量（部分可配置） ====================

my_ip: str = socket.gethostbyname(socket.gethostname())
listen_port: int = -1                        # 本机监听端口
connections: Dict[str, socket.socket] = {}   # 已建立连接的终端 {“ip:port”: socket}
message_queue: List[Dict] = []               # 接收到的消息
uid_history: List[str] = []                  # 防止消息风暴的 UID 历史记录
debug_state: bool = False                    # 调试模式(True则在控制台输出log)
max_uid_count: int = 1000                    # 最大 UID 缓存数量(防止内存超限，值为-1则不判断)

# ========================== 日志系统 ==========================

# 创建 log 文件夹
os.makedirs("log/simpleNetwork", exist_ok=True)

# 日志文件名为程序启动时间
log_file_name = f"log/simpleNetwork/{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.log"

# 写日志的底层函数
def _write_log(level: str, message: str):
    with open(log_file_name, "a", encoding="utf-8") as f:
        f.write(f"[{level}][{datetime.now().strftime('%H:%M:%S')}] {message}\n")

# 三个日志接口
def log_info(message: str):
    _write_log("INFO", message)
    if debug_state:
        print(f"[INFO] {message}")
def log_warning(message: str):
    _write_log("WARNING", message)
    if debug_state:
        print(f"\033[33m[WARNING] {message}\033[0m")
def log_error(message: str):
    _write_log("ERROR", message)
    if debug_state:
        print(f"\033[31m[ERROR] {message}\033[0m")

# ========================== 外部函数 1：设置监听端口 ==========================

def setListenPort(port: int):
    """
    设置监听端口，并启动后台线程监听其他终端连接
    """
    global listen_port
    port = int(port)
    listen_port = port

    def _listen_thread():
        server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        server.bind(("", listen_port))
        server.listen(5)
        log_info(f"监听启动：端口 {listen_port}")

        while True:
            try:
                client_socket, addr = server.accept()
                _handle_new_connection(client_socket, addr)
            except Exception as e:
                log_error(f"监听错误：{e}")

    threading.Thread(target=_listen_thread, daemon=True).start()

# ========================== 外部函数 2：连接终端 ==========================

def connect(ip: str, port: int) -> bool:
    """
    主动连接另一个终端，建立持久 socket 连接
    """
    ip = str(ip)
    port = int(port)
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.connect((ip, port))
        key = f"{ip}:{port}"
        connections[key] = s
        threading.Thread(target=_recv_thread, args=(s,), daemon=True).start()
        log_info(f"连接成功：{key}")
        
        # 发送连接状态消息
        try:
            payload = {
                "uid": f"{time.time()}-{listen_port}-connect",
                "src_ip": my_ip,
                "src_port": listen_port,
                "trg_ip": ip,
                "trg_port": port,
                "type": "connect",
                "content": ""
            }
            s.send((str(payload) + "\n").encode())
        except Exception as e:
            log_error(f"发送连接消息失败：{e}")
            
        return True
    except Exception as e:
        log_error(f"连接失败：{ip}:{port} - {e}")
        return False

# ========================== 外部函数 3：关闭连接 ==========================

def close(ip: str, port: int) -> bool:
    """
    优雅地关闭与指定终端的连接
    """
    ip = str(ip)
    port = int(port)
    
    # 如果是关闭所有连接
    if ip == "all" and port == -1:
        success = True
        # 复制当前连接列表，避免在迭代时修改字典
        connections_copy = connections.copy()
        for conn_key in connections_copy:
            conn_ip, conn_port = conn_key.split(":")
            conn_port = int(conn_port)
            # 调用单个连接关闭函数
            if not close(conn_ip, conn_port):
                success = False
        return success
    
    key = f"{ip}:{port}"
    
    if key in connections:
        try:
            # 发送关闭信号消息
            payload = {
                "uid": f"{time.time()}-{listen_port}-close",
                "src_ip": my_ip,
                "src_port": listen_port,
                "trg_ip": ip,
                "trg_port": port,
                "content": "__CLOSE_SIGNAL__",
                "type": "close"
            }
            connections[key].send((str(payload) + "\n").encode())
            
            connections[key].close()
            del connections[key]
            log_info(f"连接已关闭：{key}")
            # 添加关闭状态消息
            message_queue.append({
                "src_ip": ip,
                "src_port": port,
                "state": "close",
                "content": ""
            })
            return True
        except Exception as e:
            log_error(f"关闭连接时出错：{key} - {e}")
            return False
    else:
        log_warning(f"尝试关闭不存在的连接：{key}")
        return False

# ========================== 外部函数 4：发送消息 ==========================

def sendMessage(ip: str, port: int, content: str) -> bool:
    """
    向指定终端或广播发送消息
    """
    ip = str(ip)
    port = int(port)
    content = str(content)
    uid = f"{time.time()}-{listen_port}"  # 消息唯一 ID 防止风暴
    # 限制历史记录队列长度，防止内存无限增长
    uid_history.append(uid)
    while len(uid_history) > 1000:
        uid_history.pop(0)  # 移除最旧的记录

    # 防呆
    if (ip == "all" and port != -1) or (ip != "all" and port == -1):
        log_warning("参数错误：广播时 ip 和 port 必须同时指定")
        return False

    if ip == "all" and port == -1:
        # 广播模式
        for conn_key, sock in connections.items():
            _send(sock, content, uid, my_ip, listen_port, ip, port)
        return True

    target_key = f"{ip}:{port}"
    if target_key in connections:
        _send(connections[target_key], content, uid, my_ip, listen_port, ip, port)
        return True
    else:
        # 没有目标终端，广播给所有已连接节点转发
        for conn_key, sock in connections.items():
            _send(sock, content, uid, my_ip, listen_port, ip, port)
        return True

# ========================== 外部函数 5：获取消息 ==========================

def getMessage() -> List[Dict]:
    """
    返回接收到的所有消息，并清空缓存
    """
    global message_queue
    msgs = message_queue.copy()
    message_queue.clear()
    return msgs

# ========================== 内部函数：处理新连接 ==========================

def _handle_new_connection(sock: socket.socket, addr):
    """
    内部函数：接收到新连接，添加到 connections 并启动监听线程
    """
    key = f"{addr[0]}:{addr[1]}"
    connections[key] = sock
    threading.Thread(target=_recv_thread, args=(sock,), daemon=True).start()
    log_info(f"新连接加入：{key}")
    
    # 发送连接状态消息
    try:
        payload = {
            "uid": f"{time.time()}-{listen_port}-connect",
            "src_ip": my_ip,
            "src_port": listen_port,
            "trg_ip": addr[0],
            "trg_port": addr[1],
            "type": "connect",
            "content": ""
        }
        sock.send((str(payload) + "\n").encode())
    except Exception as e:
        log_error(f"发送连接消息失败：{e}")

# ========================== 内部函数：发送消息结构 ==========================

def _send(sock: socket.socket, content: str, uid: str, src_ip: str, src_port: int, trg_ip: str, trg_port: int):
    """
    内部函数：将消息打包为 JSON 并发送
    """
    try:
        payload = {
            "uid": uid,
            "src_ip": src_ip,
            "src_port": src_port,
            "trg_ip": trg_ip,
            "trg_port": trg_port,
            "type": "stay",
            "content": content
        }
        sock.send((str(payload) + "\n").encode())
        log_info(f"发送成功：{payload}")
    except Exception as e:
        log_error(f"发送失败：{e}")

# ========================== 内部函数：接收线程 ==========================

def _recv_thread(sock: socket.socket):
    """
    内部函数：持续接收数据并处理消息（含转发）
    """
    buffer = ""
    addr = None
    try:
        addr = sock.getpeername()
    except Exception as e:
        log_error(f"无法获取连接地址，错误信息：{e}")
        
    while True:
        try:
            data = sock.recv(1024)
            if not data:
                break
            buffer += data.decode()

            while "\n" in buffer:
                line, buffer = buffer.split("\n", 1)
                _handle_incoming_message(eval(line))  # 假定内容为字典结构字符串
        except Exception as e:
            log_error(f"接收线程出错：{e}")
            break
    
    # 连接断开时添加断开状态消息
    if addr:
        message_queue.append({
            "src_ip": addr[0],
            "src_port": addr[1],
            "state": "disconnect",
            "content": ""
        })
        key = f"{addr[0]}:{addr[1]}"
        if key in connections:
            del connections[key]

# ========================== 内部函数：处理消息 ==========================

def _handle_incoming_message(msg: dict):
    """
    内部函数：处理接收到的消息或转发
    """
    # 检查是否为关闭信号
    if msg.get("type") == "close":
        src_ip = msg.get("src_ip")
        src_port = msg.get("src_port")
        key = f"{src_ip}:{src_port}"
        
        if key in connections:
            connections[key].close()
            del connections[key]
            
        message_queue.append({
            "src_ip": src_ip,
            "src_port": src_port,
            "state": "close",
            "content": ""
        })
        return
    
    # 检查是否为连接消息
    if msg.get("type") == "connect":
        message_queue.append({
            "src_ip": msg.get("src_ip"),
            "src_port": msg.get("src_port"),
            "state": "connect",
            "content": ""
        })
        return
    
    uid = msg.get("uid")
    if uid in uid_history:  # 检查是否已处理过该UID
        return  # 防止风暴
    # 限制历史记录队列长度，防止内存无限增长
    uid_history.append(uid)
    while len(uid_history) > 1000:
        uid_history.pop(0)  # 移除最旧的记录

    trg_ip = msg.get("trg_ip")
    trg_port = msg.get("trg_port")
    content = msg.get("content")
    src_ip = msg.get("src_ip")
    src_port = msg.get("src_port")

    # 判断是不是发给自己的消息
    if (trg_ip == "all" and trg_port == -1) or ((trg_ip == my_ip or trg_ip == "127.0.0.1") and trg_port == listen_port):
        message_queue.append({
            "src_ip": src_ip,
            "src_port": src_port,
            "state": "stay",
            "content": content
        })
        log_info(f"收到消息：{src_ip}:{src_port} -> {content}")
    if (trg_ip == "all" and trg_port == -1) or not((trg_ip == my_ip or trg_ip == "127.0.0.1") and trg_port == listen_port):
        # 不是自己的消息，转发给已知连接
        for conn_key, sock in connections.items():
            _send(sock, content, uid, src_ip, src_port, trg_ip, trg_port)
        log_info(f"转发消息：{uid} -> {trg_ip}:{trg_port}")